import cv2
import numpy as np
from PIL import Image
import pytesseract
import logging

# 设置日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def preprocess_image(image_path):
    """对图像进行预处理（灰度化和二值化）"""
    try:
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像文件: {image_path}")

        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, binary = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY_INV)
        return binary
    except Exception as e:
        logging.error(f"图像预处理失败: {e}")
        raise


def detect_text_regions(binary_image):
    """自动检测文本区域"""
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    bounding_boxes = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        # 过滤掉过小的区域，并假设运单号区域有一定的宽度和高度
        if h > 20 and w > 100:
            bounding_boxes.append((x, y, w, h))

    # 按y坐标排序，确保从上到下读取
    bounding_boxes.sort(key=lambda b: b[1])
    return bounding_boxes


def crop_regions(binary_image, regions):
    """根据定义的区域裁剪图像"""
    cropped_images = []
    for region in regions:
        x, y, w, h = region
        cropped_image = binary_image[y:y + h, x:x + w]
        cropped_images.append(cropped_image)
    return cropped_images


def ocr_region(region_image, lang='chi_sim'):
    """对单个区域的图像进行OCR识别"""
    try:
        custom_config = r'--oem 3 --psm 6 -l {}'.format(lang)
        text = pytesseract.image_to_string(Image.fromarray(region_image), config=custom_config)
        return text.strip()
    except Exception as e:
        logging.error(f"区域OCR识别失败: {e}")
        raise


def ocr_core(image_path, lang='chi_sim'):
    """主函数：分区域识别图片中的文本"""
    try:
        binary_image = preprocess_image(image_path)
        regions = detect_text_regions(binary_image)

        cropped_images = crop_regions(binary_image, regions)

        results = []
        for i, cropped_image in enumerate(cropped_images):
            text = ocr_region(cropped_image, lang)
            results.append((i + 1, text))

        return results
    except Exception as e:
        logging.error(f"OCR核心处理失败: {e}")
        raise


if __name__ == "__main__":
    # 示例调用
    try:
        image_path = '/home/weiqiangren/Python/project/py-basic-exercises/图片文字识别/企业微信截图_17387351183679.png'  # 替换为你的图片路径
        result = ocr_core(image_path)
        for region_id, text in result:
            print(f"Region {region_id}: {text}")
    except Exception as e:
        logging.error(f"主程序执行失败: {e}")